import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 检查是否有可用的 GPU，如果没有则使用 CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")


# 1-定义生成器和判别器
class Generator(nn.Module):
    def __init__(self, z_dim=100, img_dim=784):
        super(Generator, self).__init__()
        self.gen = nn.Sequential(
            nn.Linear(z_dim, 256),
            nn.ReLU(),
            nn.Linear(256, img_dim),
            nn.Tanh(),  # 输出范围在 -1 到 1 之间
        )

    def forward(self, x):
        return self.gen(x)


class Discriminator(nn.Module):
    def __init__(self, img_dim=784):
        super(Discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Linear(img_dim, 128),
            nn.LeakyReLU(0.01),
            nn.Linear(128, 1),
            nn.Sigmoid(),  # 输出范围在 0 到 1 之间
        )

    def forward(self, x):
        return self.dis(x)


# 2-定义训练数据和数据加载器
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 3-初始化生成器和判别器，并将它们移动到设备上
z_dim = 100
img_dim = 28 * 28

generator = Generator(z_dim, img_dim).to(device)
discriminator = Discriminator(img_dim).to(device)

lr = 0.0001
num_epochs = 100

optimizer_gen = optim.Adam(generator.parameters(), lr=lr)
optimizer_dis = optim.Adam(discriminator.parameters(), lr=lr)

criterion = nn.BCELoss()

# 4-训练过程
for epoch in range(num_epochs):
    for batch_idx, (real, _) in enumerate(train_loader):
        real = real.view(-1, img_dim)  # 将图像展平
        batch_size = real.size(0)

        # 训练判别器
        noise = torch.randn(batch_size, z_dim, device=device)
        fake = generator(noise)
        disc_real_loss = criterion(discriminator(real), torch.ones(batch_size, 1, device=device))
        disc_fake_loss = criterion(discriminator(fake.detach()), torch.zeros(batch_size, 1, device=device))
        disc_loss = (disc_real_loss + disc_fake_loss) / 2

        optimizer_dis.zero_grad()
        disc_loss.backward()
        optimizer_dis.step()

        # 训练生成器
        noise = torch.randn(batch_size, z_dim, device=device)
        fake = generator(noise)
        gen_loss = criterion(discriminator(fake), torch.ones(batch_size, 1, device=device))

        optimizer_gen.zero_grad()
        gen_loss.backward()
        optimizer_gen.step()

        if batch_idx % 100 == 0:
            print(f"Epoch [{epoch}/{num_epochs}] Batch {batch_idx}/{len(train_loader)} \
                  Loss D: {disc_loss:.4f}, loss G: {gen_loss:.4f}")

    # 每个 epoch 保存一些生成的图像
    generator.eval()
    with torch.no_grad():
        noise = torch.randn(1, z_dim, device=device)
        generated_img = generator(noise).view(28, 28).cpu().numpy()
        plt.imshow(generated_img, cmap='gray')
        plt.savefig(f'generated_img_epoch_{epoch}.png')
        plt.close()
    generator.train()

# 5-保存模型
torch.save(generator.state_dict(), 'generator.pt')
torch.save(discriminator.state_dict(), 'discriminator.pt')
